package com.example.springai.controller;

import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

@RestController
@RequestMapping("/chroma")
public class ChromaVectorStoreController {
    private final VectorStore vectorStore;

    public ChromaVectorStoreController(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    /**
     * 向知识库添加内容
     * @param text 文本内容
     * @param storeId 知识库ID
     * @return
     */
    @RequestMapping("/add")
    public String add(String text, String storeId) {
        // 构造文档，并把归属的知识库添加到元数据中
        List<Document> docs = List.of(new Document(text, Map.of("storeId", storeId)));
        vectorStore.add(docs);
        return "Add Success";
    }

    /**
     * 知识库查询
     * @param query 查询内容
     * @param storeId 知识库ID
     * @return
     */
    @RequestMapping("/query")
    public List<Document> query(String query, String storeId) {
        FilterExpressionBuilder b = new FilterExpressionBuilder();

        List<Document> documents = vectorStore.similaritySearch(
                SearchRequest.defaults()
                        .withQuery(query)   // 查询内容
                        .withTopK(3)        // 查询数量
//                        .withSimilarityThreshold(SIMILARITY_THRESHOLD)  // 相似度阈值
                        .withFilterExpression(b.eq("storeId", storeId).build()) // 基于元数据的过滤条件
        );
        return documents;
    }
}
